Solving Multi-codebook Quantization in the GPU

We focus on the problem of vector compression using multicodebook quantization (MCQ). MCQ is a generalization of k-means where the centroids arise from the combinatorial sums of entries in multiple codebooks, and has become a critical component of large-scale, state-of-the-art approximate nearest neighbour search systems. MCQ is often addressed in an iterative manner, where learning the codebooks can be solved exactly via least-squares, but finding the optimal codes results in a large number of combinatorial NP-Hard problems. Recently, we have demonstrated that an algorithm based on stochastic local search for this problem outperforms all previous approaches. In this paper we introduce a GPU implementation of our method, which achieves a 30× speedup over a single-threaded CPU implementation. Our code is publicly available.